# -*- coding: utf-8 -*-
# @Time    : 2024/12/26 上午10:16
# @Author  : ysj
# @FileName: 4_rtsp.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/ydscc?type=blog
import cv2
import os
import numpy as np

# 定义四个 RTSP 地址和 ID


import cv2
import numpy as np

# RTSP stream URLs (replace these with your actual URLs)

rtsp_urls = [
    # "rtsp://admin:zx123456@192.168.1.63:554/h265/ch1/main/av_stream",
    # "rtsp://admin:yzlaji666@192.168.1.64:554/h265/ch1/main/av_stream",
    # "rtsp://admin:yzlaji666@192.168.1.65:554/h265/ch1/main/av_stream",
    # "rtsp://admin:yzlaji666@192.168.1.66:554/h265/ch1/main/av_stream"
]
# Initialize the video captures for each stream
# Initialize video captures for each RTSP stream
captures = [cv2.VideoCapture(url) for url in rtsp_urls]

# Define video writers to save the streams to local files
fourcc = cv2.VideoWriter_fourcc(*'XVID')  # Video codec
output_videos = [
    # cv2.VideoWriter('output_1.avi', fourcc, 30.0, (640, 480)),
    # cv2.VideoWriter('output_2.avi', fourcc, 30.0, (640, 480)),
    # cv2.VideoWriter('output_3.avi', fourcc, 30.0, (640, 480)),
    cv2.VideoWriter('output_4.avi', fourcc, 30.0, (640, 480))
]

# Frame storage for current frames
frames = [None] * 4

# Define window names for the individual streams
window_names = ["Stream 1", "Stream 2", "Stream 3", "Stream 4"]

# Target size for resizing the frames
target_size = (640, 480)  # Resize frames to 640x480

# Key for saving current frame (press 'v')
key = None

# Create a function to display the 2x2 grid
def display_grid(frames):
    # Resize frames to the target size
    resized_frames = [cv2.resize(frame, target_size) for frame in frames]

    # Combine the frames into a grid (2x2)
    top_row = np.hstack([resized_frames[0], resized_frames[1]])  # Horizontal stack for the top row
    bottom_row = np.hstack([resized_frames[2], resized_frames[3]])  # Horizontal stack for the bottom row
    grid = np.vstack([top_row, bottom_row])  # Vertical stack for the full grid

    # Resize to fit screen if necessary
    screen_width = 1920
    screen_height = 1080
    grid_resized = cv2.resize(grid, (screen_width, screen_height))

    # Show the grid
    cv2.imshow("RTSP Streams", grid_resized)

# Main loop
while True:
    ret_list = [cap.read() for cap in captures]  # Read frames from all RTSP streams

    # Check if all streams are successfully captured
    if not all(ret for ret, _ in ret_list):
        print("Error: Failed to capture from one or more streams.")
        break

    # Get frames from the ret_list
    frames = [frame if ret else np.zeros((480, 640, 3), dtype=np.uint8) for ret, frame in ret_list]

    # Draw IDs on frames (Top Left corner for example)
    for i, frame in enumerate(frames):
        cv2.putText(frame, f"ID: {i+1}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    # Display the frames in a 2x2 grid
    display_grid(frames)

    # Key press handling
    key = cv2.waitKey(1) & 0xFF

    # If 'v' is pressed, save the current frame of each stream
    if key == ord('v'):
        for i, frame in enumerate(frames):
            cv2.imwrite(f"frame_{i+1}.jpg", frame)  # Save the current frame as an image

    # Save video frames
    for i, frame in enumerate(frames):
        output_videos[i].write(frame)

    # If 'q' is pressed, exit the loop
    if key == ord('q'):
        break

# Release all resources
for cap in captures:
    cap.release()

for video in output_videos:
    video.release()

cv2.destroyAllWindows()